//! A generic hash table implementation in Rust.
//! This implementation is based on the C code provided, with modifications to ensure memory safety,
//! type safety, and idiomatic Rust usage.

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::mem;

/// A trait for custom key hashing functions.
pub trait HashFunc<K> {
    fn hash(&self, key: &K) -> u64;
}

/// A trait for custom key equality functions.
pub trait EqualFunc<K> {
    fn equal(&self, a: &K, b: &K) -> bool;
}

/// A trait for custom key and value freeing functions.
pub trait FreeFunc<K, V> {
    fn free_key(&self, key: K);
    fn free_value(&self, value: V);
}

/// A default implementation of `HashFunc` using Rust's `Hash` trait.
pub struct DefaultHashFunc;

impl<K: Hash> HashFunc<K> for DefaultHashFunc {
    fn hash(&self, key: &K) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }
}

/// A default implementation of `EqualFunc` using Rust's `PartialEq` trait.
pub struct DefaultEqualFunc;

impl<K: PartialEq> EqualFunc<K> for DefaultEqualFunc {
    fn equal(&self, a: &K, b: &K) -> bool {
        a == b
    }
}

/// A default implementation of `FreeFunc` that does nothing.
pub struct DefaultFreeFunc;

impl<K, V> FreeFunc<K, V> for DefaultFreeFunc {
    fn free_key(&self, _key: K) {}
    fn free_value(&self, _value: V) {}
}

/// A hash table entry containing a key-value pair and a pointer to the next entry.
struct HashTableEntry<K, V> {
    pair: (K, V),
    next: Option<Box<HashTableEntry<K, V>>>,
}

/// A hash table structure.
pub struct HashTable<K, V, H: HashFunc<K> = DefaultHashFunc, E: EqualFunc<K> = DefaultEqualFunc, F: FreeFunc<K, V> = DefaultFreeFunc> {
    table: Vec<Option<Box<HashTableEntry<K, V>>>>,
    table_size: usize,
    hash_func: H,
    equal_func: E,
    free_func: F,
    entries: usize,
    prime_index: usize,
}

/// A set of good hash table prime numbers.
const HASH_TABLE_PRIMES: &[usize] = &[
    193, 389, 769, 1543, 3079, 6151, 12289, 24593, 49157, 98317, 196613, 393241, 786433, 1572869,
    3145739, 6291469, 12582917, 25165843, 50331653, 100663319, 201326611, 402653189, 805306457,
    1610612741,
];

impl<K, V> HashTable<K, V, DefaultHashFunc, DefaultEqualFunc, DefaultFreeFunc> {
    /// Creates a new hash table with default hashing, equality, and freeing functions.
    pub fn new() -> Self {
        Self::with_funcs(DefaultHashFunc, DefaultEqualFunc, DefaultFreeFunc)
    }
}

impl<K, V, H: HashFunc<K>, E: EqualFunc<K>, F: FreeFunc<K, V>> HashTable<K, V, H, E, F> {
    /// Creates a new hash table with custom hashing, equality, and freeing functions.
    pub fn with_funcs(hash_func: H, equal_func: E, free_func: F) -> Self {
        let prime_index = 0;
        let table_size = HASH_TABLE_PRIMES[prime_index];
        let table = vec![None; table_size];
        Self {
            table,
            table_size,
            hash_func,
            equal_func,
            free_func,
            entries: 0,
            prime_index,
        }
    }

    /// Inserts a key-value pair into the hash table.
    pub fn insert(&mut self, key: K, value: V) -> bool {
        if (self.entries * 3) / self.table_size > 0 {
            if !self.enlarge() {
                return false;
            }
        }

        let index = (self.hash_func.hash(&key) % self.table_size as u64) as usize;

        if let Some(entry) = self.table[index].as_mut() {
            let mut current = entry;
            loop {
                if self.equal_func.equal(&current.pair.0, &key) {
                    self.free_func.free_key(mem::replace(&mut current.pair.0, key));
                    self.free_func.free_value(mem::replace(&mut current.pair.1, value));
                    return true;
                }
                if let Some(ref mut next) = current.next {
                    current = next;
                } else {
                    break;
                }
            }
            current.next = Some(Box::new(HashTableEntry {
                pair: (key, value),
                next: None,
            }));
        } else {
            self.table[index] = Some(Box::new(HashTableEntry {
                pair: (key, value),
                next: None,
            }));
        }

        self.entries += 1;
        true
    }

    /// Looks up a value by key in the hash table.
    pub fn lookup(&self, key: &K) -> Option<&V> {
        let index = (self.hash_func.hash(key) % self.table_size as u64) as usize;

        if let Some(entry) = self.table[index].as_ref() {
            let mut current = entry;
            loop {
                if self.equal_func.equal(&current.pair.0, key) {
                    return Some(&current.pair.1);
                }
                if let Some(ref next) = current.next {
                    current = next;
                } else {
                    break;
                }
            }
        }

        None
    }

    /// Removes a key-value pair from the hash table.
    pub fn remove(&mut self, key: &K) -> bool {
        let index = (self.hash_func.hash(key) % self.table_size as u64) as usize;

        if let Some(entry) = self.table[index].as_mut() {
            let mut current = entry;
            let mut prev: Option<&mut Box<HashTableEntry<K, V>>> = None;

            loop {
                if self.equal_func.equal(&current.pair.0, key) {
                    self.free_func.free_key(current.pair.0);
                    self.free_func.free_value(current.pair.1);

                    if let Some(prev) = prev {
                        prev.next = current.next.take();
                    } else {
                        self.table[index] = current.next.take();
                    }

                    self.entries -= 1;
                    return true;
                }

                if let Some(ref mut next) = current.next {
                    prev = Some(current);
                    current = next;
                } else {
                    break;
                }
            }
        }

        false
    }

    /// Returns the number of entries in the hash table.
    pub fn num_entries(&self) -> usize {
        self.entries
    }

    /// Enlarges the hash table to the next prime size.
    fn enlarge(&mut self) -> bool {
        let old_table = mem::replace(&mut self.table, vec![None; self.table_size]);
        let old_table_size = self.table_size;
        self.prime_index += 1;

        if self.prime_index < HASH_TABLE_PRIMES.len() {
            self.table_size = HASH_TABLE_PRIMES[self.prime_index];
        } else {
            self.table_size = self.entries * 10;
        }

        self.table = vec![None; self.table_size];
        self.entries = 0;

        for mut entry in old_table.into_iter().flatten() {
            while let Some(e) = entry.next.take() {
                self.insert(e.pair.0, e.pair.1);
                entry = e;
            }
            self.insert(entry.pair.0, entry.pair.1);
        }

        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_lookup() {
        let mut hash_table = HashTable::new();
        assert!(hash_table.insert(1, "one"));
        assert!(hash_table.insert(2, "two"));
        assert_eq!(hash_table.lookup(&1), Some(&"one"));
        assert_eq!(hash_table.lookup(&2), Some(&"two"));
        assert_eq!(hash_table.lookup(&3), None);
    }

    #[test]
    fn test_remove() {
        let mut hash_table = HashTable::new();
        assert!(hash_table.insert(1, "one"));
        assert!(hash_table.insert(2, "two"));
        assert!(hash_table.remove(&1));
        assert_eq!(hash_table.lookup(&1), None);
        assert_eq!(hash_table.lookup(&2), Some(&"two"));
    }

    #[test]
    fn test_num_entries() {
        let mut hash_table = HashTable::new();
        assert_eq!(hash_table.num_entries(), 0);
        assert!(hash_table.insert(1, "one"));
        assert_eq!(hash_table.num_entries(), 1);
        assert!(hash_table.insert(2, "two"));
        assert_eq!(hash_table.num_entries(), 2);
        assert!(hash_table.remove(&1));
        assert_eq!(hash_table.num_entries(), 1);
    }
}
